import torch.nn as nn
import torch
from collections import OrderedDict


class Fast_RCNN_VGG16(nn.Module):
    """
    fast r-cnn 的前置结构 VGG16实现
    """
    def __init__(self):
        """
        初始化网络结构
        """
        super(Fast_RCNN_VGG16, self).__init__()
        # VGG16
        self.Convolution_Layer = nn.Sequential(
            OrderedDict([('Conv_1', nn.Conv2d(3, 64, 3, stride=1, padding=1)),
                         ('ReLU_1', nn.ReLU(inplace=True)),
                         ('Conv_2', nn.Conv2d(64, 64, 3, stride=1, padding=1)),
                         ('ReLU_2', nn.ReLU(inplace=True)),
                         ('MaxPool_1', nn.MaxPool2d(2, stride=2)),

                         ('Conv_3', nn.Conv2d(64, 128, 3, stride=1, padding=1)),
                         ('ReLU_3', nn.ReLU(inplace=True)),
                         ('Conv_4', nn.Conv2d(128, 128, 3, stride=1, padding=1)),
                         ('ReLU_4', nn.ReLU(inplace=True)),
                         ('MaxPool_2', nn.MaxPool2d(2, stride=2)),

                         ('Conv_5', nn.Conv2d(128, 256, 3, stride=1, padding=1)),
                         ('ReLU_5', nn.ReLU(inplace=True)),
                         ('Conv_6', nn.Conv2d(256, 256, 3, stride=1, padding=1)),
                         ('ReLU_6', nn.ReLU(inplace=True)),
                         ('Conv_7', nn.Conv2d(256, 256, 3, stride=1, padding=1)),
                         ('ReLU_7', nn.ReLU(inplace=True)),
                         ('MaxPool_3', nn.MaxPool2d(2, stride=2)),

                         ('Conv_8', nn.Conv2d(256, 512, 3, stride=1, padding=1)),
                         ('ReLU_8', nn.ReLU(inplace=True)),
                         ('Conv_9', nn.Conv2d(512, 512, 3, stride=1, padding=1)),
                         ('ReLU_9', nn.ReLU(inplace=True)),
                         ('Conv_10', nn.Conv2d(512, 512, 3, stride=1, padding=1)),
                         ('ReLU_10', nn.ReLU(inplace=True)),
                         ('MaxPool_4', nn.MaxPool2d(2, stride=2)),

                         ('Conv_11', nn.Conv2d(512, 512, 3, stride=1, padding=1)),
                         ('ReLU_11', nn.ReLU(inplace=True)),
                         ('Conv_12', nn.Conv2d(512, 512, 3, stride=1, padding=1)),
                         ('ReLU_12', nn.ReLU(inplace=True)),
                         ('Conv_13', nn.Conv2d(512, 512, 3, stride=1, padding=1)),
                         ('ReLU_13', nn.ReLU(inplace=True)),
                         ('MaxPool_5', nn.MaxPool2d(2, stride=2))])
        )

    def forward(self, image):
        """
        前向过程
        :param image:
        :return:
        """
        conv_img = self.Convolution_Layer(image)
        return conv_img


# 得到当前设备
def get_device():
    """

    :return:
    """
    return torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')


device = get_device()

Fast_RCNN_Extractor = Fast_RCNN_VGG16().to(device)

if __name__ == '__main__':
    # 模拟一张图片输入
    image = torch.randn((1, 3, 500, 500), dtype=torch.float32).to(device)
    # 经过VGG16
    res = Fast_RCNN_Extractor(image)
    print(res)